import argparse
import os
import torch
import numpy as np
import logging
import random

from utils.config import get_train_cfg
from engine import TrainerWithSimpleScheduler as Trainer

torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True

SEED = 0
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
np.random.seed(SEED)
random.seed(SEED)

def setup_logging(logs_dir):
    BASIC_FORMAT = "[%(asctime)s][%(levelname)s] %(message)s"
    DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
    formatter = logging.Formatter(BASIC_FORMAT, DATE_FORMAT)
    
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(formatter)
    console_handler.setLevel(level=logging.INFO)

    logs_path = os.path.join(logs_dir, 'logs.txt')
    file_handler = logging.FileHandler(logs_path, mode='a', encoding='utf-8')
    file_handler.setFormatter(formatter)
    file_handler.setLevel(level=logging.INFO)
    
    logging.basicConfig(level=logging.INFO, handlers=[console_handler, file_handler])


def parse_args():
    parser = argparse.ArgumentParser(description="train net")
    parser.add_argument("--model", type=str, help='model name', required=True)
    parser.add_argument("--datasets", type=str, help='dataset names', required=True)
    parser.add_argument("--val-datasets", type=str, help="val dataset names", default=None)
    parser.add_argument("--gpus", type=str, help="use gpus id", default='0')
    parser.add_argument("--resume", type=bool, help="resume from last training", default=True)
    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = parse_args()
    config_path = 'configs/model/{}.yaml'.format(args.model)
    cfg = get_train_cfg(config_path)
    cfg.DATASETS.TRAIN = ["{}-train".format(dataset) for dataset in args.datasets.split(',')]
    cfg.DATASETS.TEST = ["{}-test".format(dataset) for dataset in (args.val_datasets.split(',') if args.val_datasets is not None else args.datasets.split(','))]
    cfg.OUTPUT_DIR = os.path.join(cfg.OUTPUT_DIR, "checkpoint", "{}/[train]-{}/[val]-{}".format(args.model, '&'.join(cfg.DATASETS.TRAIN), '&'.join(cfg.DATASETS.TEST)))
    cfg.MODEL.DEVICES = ["cuda:{}".format(gpu) for gpu in args.gpus.split(',')]

    os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
    setup_logging(cfg.OUTPUT_DIR)

    trainer = Trainer(cfg)
    trainer.resume_or_load(resume=args.resume)
    trainer.train()
